#include <queue>
#include <unordered_set>
#include "graph.h"
#include "node.h"
#include "./ops/add_op.h"
#include "../factory/factory.h"
#include "../factory/constant_factory.h"
#include "../factory/add_factory.h"

using std::queue;
using std::unordered_set;

Graph* Graph::graph = nullptr;
Graph::Graph() 
    : nodes_(new vector<Node*>()), 
      var_node2grad(new unordered_map<Node*, Node*>()),
      isGradientConstructed(false) {
}

Graph::~Graph(){
    if(var_node2grad) {
        delete var_node2grad;
        var_node2grad = nullptr;
    }
    if(nodes_) {
        for (int i = 0; i < nodes_ -> size(); i++) {
            auto node = (*nodes_)[i];
            if(node) delete node;
            node = nullptr;
        }
        delete nodes_;
        nodes_ = nullptr;
    }
}

void Graph::addNode(Node * node) {
    nodes_ -> emplace_back(node);
}

Graph* Graph::getGraph() {
    if(graph == nullptr)
        graph = new Graph();
    return graph;
}

void Graph::constructGradientNodeds(Node* root){
    // if the graph has been constructed, then return 
    if(isGradientConstructed) return;
    // BFS search queue
    queue<Node*> queue; 
    // the mapping from node to the grad input of this node 
    unordered_map<Node*, Node*> node2grad; 
    // the mapping from node to the num of node out
    unordered_map<Node*, unsigned int> node2num = getOutNumOfNodesMap(root);
    //init the queue
    queue.push(root);
    //init the mapping
    auto root_shape = root -> getOutShape();
    Factory* factory;
    factory = new ConstantFactory(Matrix(root_shape -> rows, root_shape -> cols, 1));
    Node* root_grad(factory -> createNode());
    delete factory;
    node2grad[root] = root_grad;
    node2num[root] = 1;
    while(!queue.empty()){
        auto top = queue.front();
        queue.pop();
        // if all differentiate ways have finished, then go ahead
        if(--node2num[top] == 0) {
            // if the top node is not opNode, then it has no child nodes and has no gradient function
            if(!top -> isOp()) continue;
            OpNode* op = dynamic_cast<OpNode*>(top);
            auto children = op -> getNodeIn();
            auto grad_nodes = op -> getGradNodes(node2grad[op]);
            for(int i = 0; i < children -> size(); i++){
                auto child = (*children)[i];
                // if the node has multiple ways to differentiate, then add all the gradients
                if(node2grad.find(child) != node2grad.end()){
                    factory = new AddFactory(node2grad[child], grad_nodes[i]);
                    Node* grad_child = factory -> createNode();
                    node2grad[child] = grad_child;
                    delete factory;
                    // store the variable grad, the root of the grad flow stream, to get all gradient nodes
                    if(child -> isVariable()) {
                        (*var_node2grad)[child] = grad_child;
                    }
                }else{
                    // store the variable grad, the root of the grad flow stream, to get all gradient nodes
                    if(child -> isVariable()) {
                        (*var_node2grad)[child] = grad_nodes[i];
                    }
                    node2grad[child] = grad_nodes[i];
                    // push the child into the queue if the child appear at first time
                }
                queue.push(child);
            } 
        }
    }
    isGradientConstructed = true;
}

unordered_map<Node*, unsigned int> Graph::getInNumOfGradNodesMap(){
    // the set of node realted to roots
    unordered_set<Node*> set;
    for(auto iter = var_node2grad -> begin(); iter != var_node2grad -> end(); iter++) {
        auto root = iter -> second;
        queue<Node*> queue;
        queue.push(root);    
        while(!queue.empty()){
            auto top = queue.front();
            queue.pop(); 
            if(set.find(top) != set.end()) 
                continue;
            set.insert(top);
            auto children = top -> getNodeIn();
            for(int i = 0; i < children -> size(); i++) {
                auto child = (*children)[i];
                queue.push(child);
            }
        }
    }

    // the mapping from node related to roots to the count fo the node input
    unordered_map<Node*, unsigned int> map;
    for(auto node : set){
        map[node] = node -> getNodeIn() -> size();
    }
    return map;
}

unordered_map<Node*, unsigned int> Graph::getInNumOfNodesMap(Node* root){
    // the set of node realted to roots
    unordered_set<Node*> set;
    queue<Node*> queue;
    queue.push(root);    
    while(!queue.empty()){
        auto top = queue.front();
        queue.pop(); 
        if(set.find(top) != set.end()) 
            continue;
        set.insert(top);
        auto children = top -> getNodeIn();
        for(int i = 0; i < children -> size(); i++) {
            auto child = (*children)[i];
            queue.push(child);
        }
    }

    // the mapping from node related to roots to the count fo the node input
    unordered_map<Node*, unsigned int> map;
    for(auto node : set){
        map[node] = node -> getNodeIn() -> size();
    }
    return map;
}

unordered_map<Node*, unsigned int> Graph::getOutNumOfNodesMap(Node* root) {
    // the set of node realted to roots
    unordered_map<Node*, unsigned int> map;
    queue<Node*> queue;
    queue.push(root);    
    while(!queue.empty()){
        auto top = queue.front();
        queue.pop(); 
        if(map.find(top) != map.end()) { 
            map[top] += 1;
            continue;
        } else {
            map[top] = 1;
        }
        auto children = top -> getNodeIn();
        for(int i = 0; i < children -> size(); i++) {
            auto child = (*children)[i];
            queue.push(child);
        }
    }

    return map;
}

unordered_map<Node*, unsigned int> Graph::getOutNumOfGradNodesMap(){
    // the set of node realted to roots
    unordered_map<Node*, unsigned int> map;
    for(auto iter = var_node2grad -> begin(); iter != var_node2grad -> end(); iter++) {
        auto root = iter -> second;
        queue<Node*> queue;
        queue.push(root);    
        while(!queue.empty()){
            auto top = queue.front();
            queue.pop(); 
            if(map.find(top) != map.end()) { 
                map[top] += 1;
                continue;
            } else {
                map[top] = 1;
            }
            auto children = top -> getNodeIn();
            for(int i = 0; i < children -> size(); i++) {
                auto child = (*children)[i];
                queue.push(child);
            }
        }
    }

    return map;
}
